import torch

from torch.amp.autocast_mode import autocast
from datasets import Dataset
import pandas as pd
from transformers import AutoTokenizer, AutoModelForCausalLM, DataCollatorForSeq2Seq, TrainingArguments, Trainer

from peft import LoraConfig, TaskType, get_peft_model

# 将JSON文件转换为CSV文件
# json_path = r"D:/codes/llm_about/self-llm/dataset/huanhuan.json"
json_path = r"D:/codes/llm_about/self-llm/zzzzz_train/Qwen15B05B/short_name_10k.json"
df = pd.read_json(json_path)
ds = Dataset.from_pandas(df)

pretrained_model = "D:/codes/nlp_about/pretrained_model/hfl_llama-3-chinese-8b"

tokenizer = AutoTokenizer.from_pretrained(pretrained_model, use_fast=False,
                                          trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token

print("tokenizer about: ", tokenizer.pad_token, tokenizer.pad_token_id, tokenizer.eos_token_id)


def process_func(example):
    MAX_LENGTH = 256  # Llama分词器会将一个中文字切分为多个token，因此需要放开一些最大长度，保证数据的完整性
    input_ids, attention_mask, labels = [], [], []
    instruction = tokenizer(
        f"<|start_header_id|>user<|end_header_id|>\n\n{example['instruction'] + ':' + example['input']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n",
        add_special_tokens=False)  # add_special_tokens 不在开头加 special_tokens
    response = tokenizer(f"{example['output']}<|eot_id|>", add_special_tokens=False)
    input_ids = instruction["input_ids"] + response["input_ids"] + [tokenizer.pad_token_id]
    attention_mask = instruction["attention_mask"] + response["attention_mask"] + [1]  # 因为eos token咱们也是要关注的所以 补充为1
    labels = [-100] * len(instruction["input_ids"]) + response["input_ids"] + [tokenizer.pad_token_id]
    if len(input_ids) > MAX_LENGTH:  # 做一个截断
        input_ids = input_ids[:MAX_LENGTH]
        attention_mask = attention_mask[:MAX_LENGTH]
        labels = labels[:MAX_LENGTH]
    return {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "labels": labels
    }


tokenized_id = ds.map(process_func, remove_columns=ds.column_names)
# print("""print(tokenized_id):""")
# print(tokenized_id)

# print("""print(tokenizer.decode(tokenized_id[0]['input_ids'])):""")
# print(tokenizer.decode(tokenized_id[0]['input_ids']))
# print("input_ids.length=", len(tokenizer.decode(tokenized_id[0]['input_ids'])))

model = AutoModelForCausalLM.from_pretrained(pretrained_model,
                                             device_map="auto", torch_dtype=torch.float16)

model.enable_input_require_grads()  # 开启梯度检查点时，要执行该方法

print("model.dtype=", model.dtype)

config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    # target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    target_modules=["q_proj", "k_proj", "v_proj"],
    inference_mode=False,  # 训练模式
    r=8,  # Lora 秩
    lora_alpha=32,  # Lora alaph，具体作用参见 Lora 原理
    lora_dropout=0.1  # Dropout 比例
)
# print("config=", config)

model = get_peft_model(model, config)
model.print_trainable_parameters()
model.enable_input_require_grads()

args = TrainingArguments(
    output_dir="./output/llama3",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    logging_steps=10,
    num_train_epochs=16,
    save_steps=10,
    learning_rate=1e-4,
    save_on_each_node=True,
    gradient_checkpointing=True,
    fp16=True,
)

with autocast("cuda"):
    trainer = Trainer(
        model=model,
        args=args,
        train_dataset=tokenized_id,
        data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, padding=True),
    )

    trainer.train()
